#!/usr/bin/python
# -*- coding: utf8 -*-
import psyco
psyco.full()

from pyinmune import modelo, agente, estado, poblacion
import random

random.seed(a=0)

EstadoSano = estado.Estado('Sano')
EstadoEnfermo = estado.Estado('Enfermo')
EstadoMuerto = estado.Estado('Muerto')


class AgentePrueba(agente.Agente):
    
    def inicializar(self):
        self.vida = 0

    def encontrarseCon(self, otro, peso):
        if self.estado == EstadoSano:
            if otro.estado == EstadoEnfermo:
                if random.random() < peso:
                    self.cambiarEstado(EstadoEnfermo)
                    #print "Me enfermé!"
                    self.vida = 0
                    
    def terminarIteracion(self):
        if not self.estado == EstadoMuerto:
            self.vida += 1
            if self.estado == EstadoEnfermo:
                if self.vida > 5:
                    self.cambiarEstado(EstadoMuerto)
                    #print "Mori!"


class GeneradorDePoblacion(object):
    
    def __init__(self):
        pass

    def generar(self):
        agentes = set()
        caracteristicas = {'edad': 30, 'sexo': True}
        for i in xrange(1000):
            agentes.add(AgentePrueba(EstadoSano, caracteristicas))

        agentes.add(AgentePrueba(EstadoEnfermo, caracteristicas))

        vinculos = set()
        for i in xrange(30000):
            print i
            a, b = random.sample(agentes, 2)
            vinculos.add(poblacion.Eje(a, b, random.random()))

        return poblacion.Poblacion(agentes, vinculos)


def main():
    poblacion = GeneradorDePoblacion().generar()
    m = modelo.Modelo(poblacion)
    for i in xrange(500):
        print i
        m.avanzarIteracion()
    print EstadoSano, len(EstadoSano)
    print EstadoEnfermo, len(EstadoEnfermo)
    print EstadoMuerto, len(EstadoMuerto)
    print

if __name__ == '__main__':
    main()
